import numpy as np
import matplotlib.pyplot as plt
import sys
import caffe
import os

def New_Net():
    Home=os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    root=str(Home).replace('\\','/')+'/mysrc/'
    deploy=root+'hybridCNN/hybridCNN_deploy.prototxt'
    caffe_model=root+'hybridCNN/hybridCNN_iter_700000.caffemodel'
    mean_file=root+'hybridCNN/hybridCNN_mean.npy'
    # labels_filename=root+'hybridCNN/match_label.txt'
    net=caffe.Net(deploy,caffe_model,caffe.TEST)
    transformer=caffe.io.Transformer({'data':net.blobs['data'].data.shape})
    transformer.set_transpose('data',(2,0,1))
    transformer.set_mean('data',np.load(mean_file).mean(1).mean(1))
    transformer.set_raw_scale('data',255)
    transformer.set_channel_swap('data',(2,1,0))
    return net,transformer

def classify(nt,tf,img_path):
    net=nt
    transformer=tf
    image = caffe.io.load_image(img_path)
    ans=[]
    probability=[]
    # transformed_image = transformer.preprocess('data', image)
    # net.blobs['data'].data[...] = transformed_image
    net.blobs['data'].data[...]=transformer.preprocess('data',image)
    output=net.forward()
    output_prob = output['prob'][0]
    for i in range(1,6):
        ans.append(output_prob.argsort()[-i])
        probability.append(output_prob[output_prob.argsort()[-i]])
    return (ans,probability)

def map_relationship(n):
    if n>=0 and n<=397:#animal
        return 0
    elif (n in [406,424,425,493,573,657,671,730]):#Building
        return 3
    elif (n in [496,615,952]) or (n>=954 and n<=958) or (n>=978 and n<=1182):#scence
        return 5
    elif (n in [731,953,960,961]):#other
        return 6
    elif (n>=398 and n<=903) or n==959 or n==977:#tool
        return 4
    elif n>=904 and n<=951:#food
        return 2
    elif n>=962 and n<=976:#plant
        return 1
    else:#other
        return 6

def decide(result):
    lst=[0,0,0,0,0,0,0]
    for i in range(5):
        # a=result[0][i]
        # print 'a=',a
        # b=map_relationship(a)
        # print 'b=',b
        # c=result[1][i]
        # print 'c=',c
        # lst[b]=lst[b]+c
        # print lst[b]
        lst[map_relationship(result[0][i])]+=result[1][i]
    # print 'wahaha'
    r=max(lst)
    ans=lst.index(r)
    return ans
    # for i in result[0]:
    #     lst[map_relationship(i)]+=result[1][result[0].index(i)]

# img_path='F:/test_jpg/restuarant2.jpg'
# nt,tf=New_Net()
# nss=classify(nt,tf,img_path)
# print nss
# boom=decide(nss)
# print boom